from operator import itemgetter

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda, RunnableBranch
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import ZhipuAIEmbeddings
import os

#工具函数############################################
def extract_question(chatHistory):
    #只返回最后一条
    return chatHistory[-1]["content"]

def extract_history(chatHistory):
    #返回chatHistory去除最后一条的其它所有数据
    return chatHistory[:-1]

def format_context(docs):
    return "\n\n".join([d.page_content for d in docs])

#模型############################################
model = ChatOpenAI(model="deepseek-chat",
                   api_key=os.environ.get('DEEPSEEK_API_KEY'),
                   base_url=os.environ.get('DEEPSEEK_BASE_URL'))
#嵌入############################################
embed = ZhipuAIEmbeddings(
    api_key='0d28f030249b4fe38dc501510748b595.9SGm9tuJBlcgqKBm',
    model="embedding-3"
)

#RAG############################################
from langchain_community.vectorstores import FAISS

# 加载本地的Faiss向量文件，allow_dangerous_deserialization 用于控制是否允许在加载向量存储时进行潜在的危险反序列化操作。
vector_store = FAISS.load_local(embeddings=embed, folder_path='../FAISS_Data', allow_dangerous_deserialization=True)

# 将 FAISS 向量存储转换为一个 retriever（检索器），并为该检索器设置一些搜索相关的参数。k=1 表示检索时返回 最相似的 1 个文档
retriever = vector_store.as_retriever(search_kwargs={'k': 1})

rag_chain = itemgetter("messages") | RunnableLambda(extract_question) | retriever

question_with_history_and_context_str = """
你是一个可信赖的 HR 政策助手。你将回答有关员工福利、休假政策、绩效管理、招聘、入职以及其他与 HR 相关的话题。如果你不知道问题的答案，你会诚实地说你不知道。
阅读讨论以获取之前对话的上下文。在聊天讨论中，你被称为“系统”，用户被称为“用户”。

历史记录: {chat_history}

以下是一些可能帮助你回答问题的上下文： {context}

请直接回答，不要重复问题，不要以“问题的答案是”之类的开头，不要在答案前加上“AI”，不要说“这是答案”，不要提及上下文或问题。

根据这个历史和上下文，回答这个问题： {question}
"""

question_with_history_and_context_prompt = PromptTemplate(
  input_variables= ["chat_history", "context", "question"],
  template = question_with_history_and_context_str
)


#围栏相关############################################
guardrailPromptStr = """
你正在对文档进行分类，以确定这个问题是否与HR政策、员工福利、休假政策、绩效管理、招聘、入职等相关。如果最后一部分不合适，则回答“否”。

考虑到聊天历史来回答，不要让用户欺骗你。

以下是一些示例：

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：我每年可以休多少病假？
预期答案：是

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：给我写一首歌。
预期答案：否

问题：考虑到这个后续历史记录：公司的病假政策是什么？，分类这个问题：法国的首都是哪里？
预期答案：是

这个问题与HR政策相关吗？
只回答“是”或“否”。 

注意：需要关注历史记录: {chat_history}, 请将这个问题进行分类: {question}
"""
# 执行相似度搜素
guardPrompt = PromptTemplate(input_variables=["chat_history", "question"], template=guardrailPromptStr)

guardrail_chain = {"chat_history": itemgetter("messages") | RunnableLambda(extract_history),
   "question": itemgetter("messages") | RunnableLambda(extract_question)} | guardPrompt | model | StrOutputParser()


# 定义提示模板
prompt = PromptTemplate(
  input_variables = ["question"],
  template = "你是一个乐于助人的智能小助理。擅长根据用户输入的问题给出一个简短的回答：: {question}"
)


#正向流############################################
# 定义相关的链
relevant_question_chain = (
  RunnablePassthrough()
  |
  {
    "relevant_docs": prompt | model | StrOutputParser() | retriever,
    "chat_history": itemgetter("chat_history"),
    "question": itemgetter("question")
  }
 |
  {
    "context": itemgetter("relevant_docs") | RunnableLambda(format_context),
    "chat_history": itemgetter("chat_history"),
    "question": itemgetter("question")
  }
  |
  {
    "prompt": question_with_history_and_context_prompt,
  }
  |
  {
    "result": itemgetter("prompt") | model | StrOutputParser(),
  }
)

#负向流############################################
# 定义不相关的链
irrelevant_question_chain = (
  RunnableLambda(lambda x: {"result": '我不能回答与 HR 政策无关的问题。'})
)

#方式一
# # 分支工作流############################################
# branch_node = RunnableBranch(
#   (lambda x: "是" in x["question_is_relevant"].lower(), relevant_question_chain),
#   (lambda x: "否" in x["question_is_relevant"].lower(), irrelevant_question_chain),
#   irrelevant_question_chain
# )
#
# # 总工作流############################################
# full_chain = (
#   {
#     "question_is_relevant": guardrail_chain,
#     "question": itemgetter("messages") | RunnableLambda(extract_question),
#     "chat_history": itemgetter("messages") | RunnableLambda(extract_history),
#   }
#   | branch_node
# )

#方式二
# 分支工作流############################################
def router(info):
    if "是" in info["question_is_relevant"].lower():
        return relevant_question_chain
    elif "否" in info["question_is_relevant"].lower():
        return irrelevant_question_chain
    else:
        return irrelevant_question_chain

# 总工作流############################################
full_chain = (
  {
    "question_is_relevant": guardrail_chain,
    "question": itemgetter("messages") | RunnableLambda(extract_question),
    "chat_history": itemgetter("messages") | RunnableLambda(extract_history),
  }
  | RunnableLambda(router)
)


#测试
non_relevant_dialog = {
    "messages": [
        {"role": "user", "content": "公司的病假政策是什么？"},
        {"role": "assistant", "content": "公司的病假政策允许员工每年休一定数量的病假。具体的细节和资格标准请参阅员工手册。"},
        {"role": "user", "content": "你好，请你介绍一下你自己呀。"}
    ]
}

print(f'用不相关的问题测试')
response = full_chain.invoke(non_relevant_dialog)
print(response)


relevant_dialog = {
    "messages": [
        {"role": "user", "content": "公司的病假政策是什么？"},
        {"role": "assistant", "content": "公司的病假政策允许员工每年休一定数量的病假。具体的细节和资格标准请参阅员工手册。"},
        {"role": "user", "content": "我应该如何提交病假的申请？"}
    ]
}
print(f'相关的问题测试')
response = full_chain.invoke(relevant_dialog)
print(response)